package com.idanchuang.component.mq.amqp.pressure;

import com.alibaba.fastjson.JSON;
import com.idanchuang.component.core.helper.EnvContext;
import com.idanchuang.component.logback.helper.LogContext;
import com.idanchuang.component.mq.amqp.processor.BeforePublishPostProcessor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.amqp.AmqpException;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.rabbit.core.RabbitTemplate;

import java.net.URLEncoder;

/**
 * @author yjy
 * Created at 2021/11/9 3:24 下午
 */
public class PressureBeforePublishPostProcessor implements BeforePublishPostProcessor {

    private static final Logger rabbitTemplateLogger = LoggerFactory.getLogger(RabbitTemplate.class);

    @Override
    public Message postProcessMessage(Message message) throws AmqpException {
        try {
            // 压测标
            String headerName = EnvContext.getEnvHeaderName();
            String headerValue = EnvContext.getCurrentEnv();
            rabbitTemplateLogger.debug("BeforePublishPostProcessors add header: {}, value: {}", headerName, headerValue);
            if (headerValue != null) {
                //拦截逻辑就是如果是压测请求就加个header标记
                message.getMessageProperties().getHeaders().put(headerName, headerValue);
            }
            // 如果有自定义链路id, 则传递
            String dcTraceIdValue = EnvContext.getDcTraceId();
            if (dcTraceIdValue != null) {
                String dcTraceIdName = EnvContext.getDcTraceIdHeaderName();
                rabbitTemplateLogger.debug("BeforePublishPostProcessors add header: {}, value: {}", dcTraceIdName, dcTraceIdValue);
                //拦截逻辑就是如果是压测请求就加个header标记
                message.getMessageProperties().getHeaders().put(dcTraceIdName, dcTraceIdValue);
            }
            // 如果有环境变量, 则传递
            if (EnvContext.getAttributeCount() > 0) {
                String attributes = JSON.toJSONString(EnvContext.getAttributes());
                attributes = URLEncoder.encode(attributes, "UTF-8");
                message.getMessageProperties().getHeaders().put(EnvContext.TRANSMIT_ATTRIBUTE_HEADER_NAME, attributes);
                rabbitTemplateLogger.debug("BeforePublishPostProcessors add header: {}, value: {}", EnvContext.TRANSMIT_ATTRIBUTE_HEADER_NAME, attributes);
            }
        } catch (Exception e) {
            // 降级
            LogContext.setKeyword("componentError");
            rabbitTemplateLogger.error("Amqp setBeforePublishPostProcessors error", e);
        }
        return message;
    }
}
